"""Minimal reproduction of the PyAV-backend keyframe-snap bug.

Synthesises a 200-frame H.264 clip with one keyframe (gop_size=200, no
B-frames) and a known signal: each frame N is a solid color whose green
channel equals N. Decodes 8 uniformly-spaced frames via the actual vLLM
loader (`VideoBackend.load_bytes`) for both backends, then recovers
which frame each backend returned by reading the green channel —
independent of the metadata label the loader claims.

Pre-fix: pyav returns 8 copies of frame 0, all labelled as distinct
target indices. Post-fix: pyav matches opencv pixel-for-pixel.
"""

import io
import sys
import types

# `vllm.platforms.cuda` (loaded by vllm/__init__.py) eagerly imports
# vllm._C_stable_libtorch for op-registration side effects. Some local
# editable builds don't include that .so file, but no code reads
# attributes from it. Stubbing lets us import VideoBackend without a
# full rebuild. Drop this hack if your build already provides the
# extension.
sys.modules.setdefault(
    "vllm._C_stable_libtorch", types.ModuleType("vllm._C_stable_libtorch")
)

import av  # noqa: E402
import numpy as np  # noqa: E402

from vllm.multimodal.video import VideoBackend  # noqa: E402

NUM_FRAMES = 200
FPS = 30
WIDTH, HEIGHT = 64, 64
NUM_SAMPLED = 8


def synthesise_long_gop_video() -> bytes:
    """Encode 200 frames of solid-color (green = frame index) with GOP=200."""
    buf = io.BytesIO()
    container = av.open(buf, mode="w", format="mp4")
    stream = container.add_stream("h264", rate=FPS)
    stream.width = WIDTH
    stream.height = HEIGHT
    stream.pix_fmt = "yuv420p"
    stream.codec_context.gop_size = NUM_FRAMES   # one keyframe, at frame 0
    stream.codec_context.max_b_frames = 0        # PTS == decode order
    stream.codec_context.options = {
        "x264-params": "scenecut=0:keyint=200:min-keyint=200"
    }

    for i in range(NUM_FRAMES):
        img = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
        img[:, :, 1] = i                         # green channel encodes index
        frame = av.VideoFrame.from_ndarray(img, format="rgb24")
        for packet in stream.encode(frame):
            container.mux(packet)
    for packet in stream.encode():
        container.mux(packet)
    container.close()
    return buf.getvalue()


def count_keyframes(data: bytes) -> int:
    with av.open(io.BytesIO(data)) as c:
        s = c.streams.video[0]
        return sum(1 for p in c.demux(s) if p.is_keyframe)


def main() -> None:
    data = synthesise_long_gop_video()
    n_keyframes = count_keyframes(data)
    print(f"Encoded {NUM_FRAMES} frames, {len(data)} bytes, "
          f"{n_keyframes} keyframe(s).")
    assert n_keyframes == 1, (
        f"Test fixture is invalid: expected 1 keyframe, got {n_keyframes}. "
        "The bug is only visible when sampled frames span a long GOP."
    )

    opencv_frames, opencv_meta = VideoBackend.load_bytes(
        data, num_frames=NUM_SAMPLED, backend="opencv"
    )
    pyav_frames, pyav_meta = VideoBackend.load_bytes(
        data, num_frames=NUM_SAMPLED, backend="pyav"
    )

    requested = list(opencv_meta["frames_indices"])
    assert requested == list(pyav_meta["frames_indices"]), (
        f"Backends disagree on metadata indices "
        f"({opencv_meta['frames_indices']} vs {pyav_meta['frames_indices']}) "
        "— separate problem."
    )
    print(f"Both backends report frames_indices: {requested}")

    # Recover the frame each backend ACTUALLY returned by reading the green
    # channel of the centre pixel. Independent of the metadata label.
    opencv_actual = [int(f[HEIGHT // 2, WIDTH // 2, 1]) for f in opencv_frames]
    pyav_actual = [int(f[HEIGHT // 2, WIDTH // 2, 1]) for f in pyav_frames]

    print()
    print(f"{'slot':>4} {'requested':>10} {'opencv got':>11} {'pyav got':>9} "
          f"{'mean|Δ|':>9}")
    print("-" * 50)
    for i, idx in enumerate(requested):
        diff = float(np.mean(np.abs(
            opencv_frames[i].astype(int) - pyav_frames[i].astype(int))))
        print(f"{i:>4} {idx:>10} {opencv_actual[i]:>11} "
              f"{pyav_actual[i]:>9} {diff:>9.2f}")

    print()
    opencv_unique = len(set(opencv_actual))
    pyav_unique = len(set(pyav_actual))
    print(f"opencv returned {opencv_unique} distinct frames")
    print(f"pyav   returned {pyav_unique} distinct frames")
    if opencv_unique == pyav_unique == NUM_SAMPLED:
        print("\nPASS: pyav and opencv agree, all requested frames returned.")
    elif pyav_unique < opencv_unique:
        print(f"\nFAIL: pyav collapsed {NUM_SAMPLED} requested slots onto "
              f"{pyav_unique} distinct frame(s) — keyframe-snap bug present.")
    else:
        print(f"\nFAIL: backends disagree (opencv={opencv_unique}, "
              f"pyav={pyav_unique} distinct frames).")


if __name__ == "__main__":
    main()